package com.shujia.flink.state;

import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.HashMap;

public class Demo3ValueState {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(2);

        DataStream<String> wordsDS = env.socketTextStream("master", 8888);

        //安装单词分组
        KeyedStream<String, String> keyByDS = wordsDS.keyBy(word -> word);

        DataStream<Tuple2<String, Integer>> countDS = keyByDS
                .process(new KeyedProcessFunction<String, String, Tuple2<String, Integer>>() {

                    /**
                     * flink的状态，再内部会为每一个key都保存一个值
                     * 状态会被checkpoint定期持久化到hdfs中
                     */
                    ValueState<Integer> valueState;

                    //open方法每一个task启动的时候执行一次，一般用于初始化
                    //需要再open方法初始化状态
                    @Override
                    public void open(Configuration parameters) throws Exception {
                        //获取flink的执行上下文对象，使用上下文对象初始化状态
                        RuntimeContext context = getRuntimeContext();
                        //创建状态描述对象，描述状态的类型和名称
                        ValueStateDescriptor<Integer> stateDescriptor = new ValueStateDescriptor<>("count", Types.INT);
                        //获取状态
                        valueState = context.getState(stateDescriptor);
                    }

                    @Override
                    public void processElement(String word,
                                               KeyedProcessFunction<String, String, Tuple2<String, Integer>>.Context ctx,
                                               Collector<Tuple2<String, Integer>> out) throws Exception {
                        //从状态中获取单词的数量
                        Integer count = valueState.value();
                        if (count == null) {
                            count = 0;
                        }
                        //累加计算单词的数量
                        count++;
                        //更新状态
                        valueState.update(count);
                        //将计算结果发送到下游
                        out.collect(Tuple2.of(word, count));
                    }
                });

        countDS.print();

        env.execute();
    }
}
